# AI分析服务模块
import ollama
import pandas as pd

class AnalysisService:
    def __init__(self, db_session, stock_api):
        self.db_session = db_session
        self.stock_api = stock_api
        self.model_name = "llama2"  # 默认模型名称
        
    def set_model(self, model_name):
        """
        设置Ollama模型
        :param model_name: 模型名称
        """
        self.model_name = model_name
        
    def get_stock_analysis(self, stock_code, days=30):
        """
        获取股票分析报告
        :param stock_code: 股票代码
        :param days: 历史数据天数
        :return: 分析报告
        """
        # 获取历史数据
        historical_data = self.stock_api.get_historical_data(stock_code, days)
        
        # 准备提示
        prompt = f"分析以下股票历史数据并提供投资建议:\n{historical_data}"
        
        # 调用Ollama模型
        response = ollama.chat(model=self.model_name, messages=[
            {"role": "user", "content": prompt}
        ])
        
        return {
            'stock_code': stock_code,
            'analysis': response['message']['content'],
            'timestamp': pd.Timestamp.now()
        }
        
    def predict_stock_trend(self, stock_code, days=7):
        """
        预测股票走势
        :param stock_code: 股票代码
        :param days: 预测天数
        :return: 预测结果
        """
        # 获取历史数据
        historical_data = self.stock_api.get_historical_data(stock_code, 90)  # 使用更多历史数据
        
        # 准备提示
        prompt = f"基于以下历史数据，预测未来{days}天的股票走势:\n{historical_data}"
        
        # 调用Ollama模型
        response = ollama.chat(model=self.model_name, messages=[
            {"role": "user", "content": prompt}
        ])
        
        return {
            'stock_code': stock_code,
            'prediction': response['message']['content'],
            'prediction_days': days,
            'timestamp': pd.Timestamp.now()
        }
        
    def generate_trading_strategy(self, user_id, risk_level="medium"):
        """
        生成交易策略
        :param user_id: 用户ID
        :param risk_level: 风险等级 (low, medium, high)
        :return: 交易策略
        """
        # 获取用户投资组合
        # 获取市场趋势
        # 生成策略
        
        return {
            'user_id': user_id,
            'risk_level': risk_level,
            'strategy': '这里是生成的交易策略',
            'timestamp': pd.Timestamp.now()
        }